# -*- coding: utf-8 -*-
import itertools
from globalvariables import *
import logging
from Agents.agent import Agent
import re
import json
from utils import *
from Agents.retrievalagent import RetrievalAgent
import asyncio
import aiohttp
import time
import copy
from itertools import islice


class ExtractionAgent(Agent):
    def __init__(self, token):
        super().__init__(token)
        self.stopwords = ["in", "of", "on", "at", "to", "with", "and", "or", "for", "a", "the", "an", "as", "via",
                          "Method", "Algorithm", "Technologies", "Function",
                          "Approach", "Process", "Architecture", 'Task',
                          "Framework", "Frameworks", "Capabilities", "Capability", "Capacity",
                          "Power", "Performance", "Improving", "Improve", "Ability", "Abilities", "Efficiency",
                          "Enhancing", "Boosting", "Enhance", "Boost", "Techniques", "Technique", "Paradigm",
                          "Approach", "Mitigating", "Mitigate", "Alleviating", "Alleviate", "Reducing", "Reduce",
                          "Phenomena", "Addressing", "Address", "Curb", "Curbing"]

    async def ExtractLanguage(self, input: str):
        """
        :param input:
        :return: survey_lang
        从用户输入中提取用户期望的语种
        """
        start = time.time()
        extr_lang_input = extract_lang_kw_prompt % input
        json_format = "{\"expected_lang\":\"the_language_expected_by_the_user\"}"
        res = await self.chat_async(extr_lang_input, do_sample=False, response_format_type="json_object",
                                    json_format=json_format)
        try:
            survey_lang = try_parse_json_object(res)[1]["expected_lang"]
        except:
            survey_lang = "English"
        end = time.time()
        print("extract_language_time:", end - start)
        return survey_lang

    def ExtractTopic(self, input: str, survey_lang: str, survey_type: int):
        """
        :param input:
        :return:
        从用户输入中提取主题.
        """
        extract_topic_prompt = extract_topic_prompts[survey_type][survey_lang]
        samples = extract_topic_samples[survey_type][survey_lang]
        extr_topic_entity_input = extract_topic_prompt % (samples, input)
        json_format = "{\"Output\":\"output\"}"
        res = self.chat(extr_topic_entity_input, response_format_type="json_object", language=survey_lang,
                        json_format=json_format, do_sample=False)
        topic = try_parse_json_object(res)[1]["Output"]
        return topic

    async def ExtractAcadQueries(self, input: str):
        """
        :param input:
        :return:
        从用户输入中提取学术查询词，第一步从用户输入中提取学术实体，第二步将学术实体转变为学术查询词”
        """
        # 提取学术实体
        # 设置样例
        start = time.time()
        samples = """input: \"我想了解更多关于深度学习中的损失函数，你能帮我做个综述吗？\", output: {\"academic_entity_1\": \"Loss Function\", \"academic_entity_2\": \"Deep Learning\"},
                     input: \"我听说Text2SQL最近挺火的，你能告诉我它目前的研究进展和面临的挑战吗？\", output: {\"academic_entity_1\": \"Text2SQL\"},
                     input: \"我想提升我的大模型的规划能力，你能帮我比较一下有哪些方法吗，它们各自的优缺点是什么？\", output: {\"academic_entity_1\": \"Planning Ability\", \"academic_entity_2\": \"Large Model\"},
                     input: \"我对多模态大模型的发展很感兴趣，你能给我讲讲它的技术发展路线吗？\", output: {\"academic_entity_1\": \"Multimodal Large Model\"}
                     """
        extr_topic_entity_input = extract_entity_prompt % (samples, input)
        json_format = """{\"academic_entity\":\"academic_entity_in_the_input\", ...} """
        response_text = await self.chat_async(extr_topic_entity_input, response_format_type="json_object",
                                              language="English",
                                              json_format=json_format)
        # 解析响应
        response_dict = try_parse_json_object(response_text)[1]
        acad_entities = []
        for key, value in response_dict.items():
            if key.startswith("academic_entity"):
                acad_entities.append(value)

        # 学术实体转变为学术查询词
        acad_queries = []
        json_format_opti = """{\"Output\":\"optimized_academic_query\", \"Explanation\":\"explanation\"}"""
        for acad_entity in acad_entities:
            need_opt = False  # 用于标记是否需要优化
            for stopword in self.stopwords:
                # 若学术实体中不含停用词，则认为不需要优化
                if stopword in acad_entity.lower():
                    need_opt = True
                    optimize_acad_query_input = optimize_acad_query_prompt % acad_entity
                    res = await self.chat_async(optimize_acad_query_input, language="English",
                                                response_format_type="json_object",
                                                json_format=json_format_opti, do_sample=False)
                    opti_q = try_parse_json_object(res)[1]["Output"]
                    if not opti_q == "None":
                        # 如果响应为None表示该词不应该转变为学术查询词，跳过
                        acad_queries.append(opti_q)
                    break
            if not need_opt:
                acad_queries.append(acad_entity)
        end = time.time()
        print("extract_entities_time:", end - start)
        # 返回学术查询词
        return acad_queries

    async def ExtractType(self, input: str):
        """
        :param input:
        :return type:
        根据用户输入划分综述类型
        """
        start = time.time()
        json_survey_type = """{\"survey_type\":\"the_type_of_the_survey\"}"""
        samples = """[{\"Input\": \"损失函数\", \"survey_type\": \"技术概念\"},
                      {\"Input\": \"Text2SQL研究现状如何，面临哪些挑战？\", \"survey_type\": \"方向研究现状\"},
                      {\"Input\": \"有哪些方法可以提升大模型的规划能力，各自优劣是什么？\", \"survey_type\": \"方法对比分析\"},
                      {\"Input\"" \"多模态大模型的技术发展路线是什么样的？\", \"survey_type\":\"技术方法发展历程\"}]"""
        classification_input = extract_type_prompt % (samples, input)
        while True:
            try:
                res = await self.chat_async(classification_input, language="Chinese",
                                            response_format_type="json_object",
                                            json_format=json_survey_type, do_sample=False)
                survey_type = try_parse_json_object(res)[1]["survey_type"]
                if survey_type in ["技术概念", "方向研究现状", "方法对比分析", "技术方法发展历程"]:
                    break
                else:
                    raise ValueError("Invalid survey type classification response")
            except ValueError as e:
                print(e)
                classification_input = classification_input + "输出应准确对应以下类别之一： {\"技术概念\", \"方向研究现状\", \"方法对比分析\", \"技术方法发展历程\"} ！"
        # 创建一个字典来映射字符串到整数
        mapping = {
            "技术概念": 0,
            "方向研究现状": 1,
            "方法对比分析": 2,
            "技术方法发展历程": 3
        }
        survey_number = mapping[survey_type]
        end = time.time()
        print(f"extract_type_time: {end - start}")
        return survey_number

    async def ExtractKeyInfo(self, chunks, survey_type: int, survey_lang: str, survey_topic: str, semaphore, model_code):
        async def process_single_chunk(chunk_text, survey_type, survey_lang, survey_topic, semaphore, i, model_code):
            start = time.time()
            extract_keyinfo_input_semi = extract_key_entity_prompts_v2_zch[survey_type][survey_lang] % chunk_text
            extract_keyinfo_input = extract_keyinfo_input_semi.replace("{综述主题}", survey_topic)
            json_format = """[{\"entity_category\": category_of_the_1st_entity, \"entity_name\": name_of_the_1st_entity, \"descriptive_linkage\": the_linkage_between_1st_entity_and_surve_topic}, {\"entity_category\": category_of_the_2nd_entity, \"entity_name\": name_of_the_2nd_entity, \"descriptive_linkage\": the_linkage_between_2nd_entity_and_surve_topic},..., {\"entity_category\": category_of_the_nth_entity, \"entity_name\": name_of_the_nth_entity, \"descriptive_linkage\": the_linkage_between_nth_entity_and_surve_topic},]"""
            async with semaphore:
                print(f"开始提取第 {i} 个chunk")
                # 异步调用 chat_async 获取返回结果
                resp = await self.chat_async(
                    user_input=extract_keyinfo_input,
                    language=survey_lang,
                    response_format_type="json_object",
                    json_format=json_format,
                    do_sample=False,  # 不启用采样，确保输出是确定性的
                    max_tokens=4095,
                    model_code=model_code
                )
                try:
                    entities = try_parse_json_object(resp)[1]
                    # 检查返回的 entities 类型，如果是字典，就将其包装成列表
                    if isinstance(entities, dict):
                        entities = [entities]  # 将字典转换为列表，或根据需要处理
                except Exception as e:
                    print(e)
                    print("Invalid JSON format")
                    return []
                end = time.time()
                print(f"提取第 {i} 个chunk用时: {end - start}s")
            return entities

        tasks = [process_single_chunk(chunk["chunk_text"], survey_type, survey_lang, survey_topic, semaphore, i, model_code)
                 for i, chunk in enumerate(chunks)]
        chunks_keyinfo = await asyncio.gather(*tasks)
        for chunk_index, chunk_keyinfo in enumerate(chunks_keyinfo):
            conference_name, _ = ext_name_year(chunks[chunk_index]["original_filename"])
            src_chunk_info = {"paper_title": chunks[chunk_index]["paper_title"],
                              "chunk_id": chunks[chunk_index]["chunk_id"],
                              "conference_name": conference_name,
                              "year": chunks[chunk_index]["year"]}
            for i in range(len(chunk_keyinfo)):
                chunks_keyinfo[chunk_index][i]["src_chunk_info"] = src_chunk_info
        # 返回的结果做一个规范化，将空列表删除，并将chunks_keyinfo元素下的字典没有'entity_category', 'entity_name', 'descriptive_linkage'三个字符的元素删掉
        for sublist in chunks_keyinfo:
            sublist[:] = [item for item in sublist if
                          all(key in item for key in ['entity_category', 'entity_name', 'descriptive_linkage'])]
        chunks_keyinfo_2 = [sublist for sublist in chunks_keyinfo if sublist]
        # 将经过规范化处理的结果返回
        return chunks_keyinfo_2

    # """异步提取单个关键信息"""
    async def extract_single_keyinfo(self, keyinfo_name, extract_keyinfo_prompt, paper_text, survey_type,
                                     survey_topic, survey_lang):
        """异步提取单个关键信息"""
        extract_keyinfo_input_semi = extract_keyinfo_prompt % paper_text

        extract_keyinfo_input = extract_keyinfo_input_semi.replace(
            survey_type_int2str[survey_lang][survey_type], survey_topic
        )
        result = ''
        for attempt in range(4):  # 重试4 次
            try:
                # 异步调用 chat_async 获取返回结果
                result = await self.chat_async(
                    extract_keyinfo_input,
                    language=survey_lang,
                    response_format_type="text",  # 设置为文本格式
                    do_sample=False  # 不启用采样，确保输出是确定性的
                )
                # 确认是否获得了有效的 result
                if result:
                    return keyinfo_name, result
                else:
                    print(f"Attempt {attempt + 1}: No result returned.")
            except Exception as e:
                # 捕获可能的异常并打印错误信息
                print(f"Error during chat_async call (Attempt {attempt + 1}): {e}")
            await asyncio.sleep(1)  # 每次重试等待 1 秒

        return keyinfo_name, "No valid result after 4 attempts"

    # """处理单篇文献，并发提取3个部分"""
    async def process_single_paper(self, chunked_paper, survey_type, survey_lang, survey_topic,
                                   extract_keyinfo_prompts):
        """处理单篇文献，并发提取3个部分"""
        # 拼接论文文本
        paper_text = "".join(chunked_paper["chunk_texts"])
        # 初始化论文关键信息
        paper_keyinfo = {
            "paper_id": chunked_paper["paper_id"],
            "paper_title": chunked_paper["paper_title"],
        }
        # 获取提示词
        extract_keyinfo_prompt_infos = extract_keyinfo_prompts[survey_type][survey_lang]
        # 并发提取关键信息
        tasks = [
            self.extract_single_keyinfo(
                keyinfo_name, extract_keyinfo_prompt, paper_text, survey_type, survey_topic, survey_lang
            )
            for keyinfo_name, extract_keyinfo_prompt in extract_keyinfo_prompt_infos.items()
        ]
        results = await asyncio.gather(*tasks)
        # 处理结果
        for keyinfo_name, res in results:
            paper_keyinfo[keyinfo_name] = res
        return paper_keyinfo

    async def extract_keyinfo_batch(self, chunked_papers_batch, survey_type, survey_lang, survey_topic,
                                    extract_keyinfo_prompts):
        """处理一批文献"""
        tasks = [
            self.process_single_paper(chunked_paper, survey_type, survey_lang, survey_topic, extract_keyinfo_prompts)
            for chunked_paper in chunked_papers_batch
        ]

        return await asyncio.gather(*tasks)

    async def ExtractKeyInfo1(self, papers_text, survey_type: int, survey_lang: str, survey_topic: str):
        # """异步处理所有文献"""
        def chunked_iterable(iterable, size):
            it = iter(iterable)
            while chunk := list(islice(it, size)):
                yield chunk

        all_keyinfo = []
        # 限制并发请求数量为一个批次处理15篇
        semaphore = asyncio.Semaphore(1)  # 限制同时只有一个批次在运行

        async def process_with_semaphore(batch):
            async with semaphore:
                return await self.extract_keyinfo_batch(batch, survey_type, survey_lang, survey_topic,
                                                        extract_keyinfo_prompts)

        for chunked_papers_batch in chunked_iterable(papers_text, 15):
            batch_results = await process_with_semaphore(chunked_papers_batch)
            all_keyinfo.extend(batch_results)
            await asyncio.sleep(0.5)  # 控制每批次之间的时间间隔，避免触发并发限制
        return all_keyinfo

    async def ExtractClaim(self, ref_chunks, survey_lang, survey_topic):
        async def ext_chunk_claims(ref_chunk, survey_lang, survey_topic):
            ext_claim_input = extract_claims_prompts[survey_lang] % (survey_topic, ref_chunk['chunk_content'])
            try:
                # 异步调用 chat
                chunk_claims = await self.chat_async(
                    ext_claim_input,
                    language=survey_lang,
                    response_format_type="text",
                    do_sample=False
                )
                ref_chunk["chunk_claims"] = chunk_claims.split('\n')
            except Exception as e:
                ref_chunk["chunk_claims"] = []

        async def process_chunk(semaphore, ref_chunk):
            async with semaphore:
                await ext_chunk_claims(ref_chunk, survey_lang, survey_topic)

        semaphore = asyncio.Semaphore(45)

        tasks = [process_chunk(semaphore, ref_chunks[i]) for i in range(len(ref_chunks))]

        # 等待所有任务完成
        await asyncio.gather(*tasks)

    # {"input": "我正在比较不同的深度学习框架，你能帮我分析一下它们在大规模数据处理上的表现吗？", "output": 3},
    # {"input": "我对强化学习的发展历史很感兴趣，你能帮我梳理一下它的技术演进脉络吗？", "output": 4},
    # {"input": "我想知道什么是大语言模型的微调，你能帮我做个综述吗？", "output": 1},
    # {"input": "我正在研究自动驾驶技术，你能告诉我计算机视觉在这方面的应用现状吗？", "output": 2},
    # {"input": "我在研究推荐系统，你能帮我分析一下不同深度学习模型在推荐系统上的表现吗？", "output": 3},
    # {"input": "我对问答系统技术的发展很感兴趣，你能帮我梳理一下它的技术发展脉络吗？", "output": 4}
    # for _ in range(10):
    #     print(a.ExtractEntity("损失函数"))
    # for _ in range(10):
    #     print(a.ExtractEntity("Text2SQL研究现状如何，面临哪些挑战？"))
    # for _ in range(10):
    #     print(a.ExtractEntity("有哪些方法可以提升大模型的规划能力，各自优劣是什么？"))
    # print(a.ExtractType("多模态大模型的技术发展路线是什么样的？"))
    # print(a.ExtractType("损失函数"))

    # print(a.ExtractType("Text2SQL研究现状如何，面临哪些挑战？"))
    # print(a.ExtractType("我对强化学习的发展历史很感兴趣，你能帮我梳理一下它的技术演进脉络吗？"))
    # print(a.ExtractType("我想知道什么是大语言模型的微调，你能帮我做个综述吗？"))
    #
    # print(a.ExtractType("我正在研究自动驾驶技术，你能告诉我计算机视觉在这方面的应用现状吗？"))
    # print(a.ExtractType("我在研究推荐系统，你能帮我分析一下不同深度学习模型在推荐系统上的表现吗？"))
    # print(a.ExtractType("我对问答系统技术的发展很感兴趣，你能帮我梳理一下它的技术发展脉络吗？"))
    #
    # print(a.ExtractType("我想了解一下基于深度学习的推荐系统，请你写一篇综述介绍介绍"))
    # print(a.ExtractType("我想了解一下强化微调，请你写一篇综述介绍介绍"))
    # print(a.ExtractType("我想了解一下迁移学习的发展历史吗，请你写一篇综述介绍介绍"))
    # print(a.ExtractType("隐式篇章关系识别任务研究现状如何？"))
    # print(a.ExtractType("事件因果关系识别任务研究现状如何？"))
    # print(a.ExtractType("不同类型的深度学习算法（如卷积神经网络、循环神经网络、强化学习算法等）在性能和应用上有哪些差异？"))
    # print(a.ExtractType("人工智能算法在数据处理和分析方面的优势和局限性是什么？"))
    # print(a.ExtractType("计算机视觉技术在图像识别和处理领域的应用现状如何？"))
    # print(a.ExtractAcadQueries("如何提高计算机视觉的精度和效率？"))
    # print(a.ExtractAcadQueries("计算机视觉技术在智能安防、医疗影像、无人驾驶等领域的应用前景和挑战有哪些？"))
    # print(a.ExtractAcadQueries("自然语言处理技术在语言理解和生成方面的发展现状和趋势是什么？"))
    # print(a.ExtractAcadQueries("如何提升自然语言处理的准确性和效率？"))
    # print(a.ExtractAcadQueries("自然语言处理技术在智能客服、文本分析、信息检索等领域的应用和创新点有哪些？"))
    # print(a.ExtractAcadQueries("机器学习算法在数据挖掘中的应用和实践情况如何？"))
    # print(a.ExtractAcadQueries("如何优化机器学习模型的性能和效果？"))
    # print(a.ExtractAcadQueries("机器学习技术在数据处理和分析方面的优势和局限性是什么？"))
    # print(a.ExtractAcadQueries("智能机器人的技术发展现状和趋势是什么？"))
    # print(a.ExtractAcadQueries("智能机器人在工业生产、服务领域、家庭生活等方面的应用和创新点有哪些？"))
    # print(a.ExtractAcadQueries("智能机器人的人机交互技术和自主决策能力如何实现？"))
    # print(a.ExtractAcadQueries("人工智能芯片的发展现状和趋势是什么？"))
    # print(a.ExtractAcadQueries("不同类型的人工智能芯片（如GPU、FPGA、ASIC等）在性能和应用上有哪些差异？"))
    # print(a.ExtractAcadQueries("人工智能芯片在推动人工智能技术发展和应用方面的作用和挑战是什么？"))
    # print(a.ExtractAcadQueries("人工智能技术在物联网领域的应用和发展情况如何？"))
    # print(a.ExtractAcadQueries("如何实现物联网设备的智能化和自动化？"))
    # print(a.ExtractAcadQueries("人工智能与物联网的融合发展带来哪些机遇和挑战？"))
